package com.ctd.utils.jes;

import cn.hutool.core.lang.Assert;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.KeyPair;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.util.Map;


/**
 * 用途：
 * 与合作伙伴 交换报文的 简洁版加解密工具
 * <p>
 * 使用方法：
 * 1. 生成sm2密钥对
 * 2. 发送者用接收者公钥加密报文
 * 3. 接收者用自己的私钥解密报文
 *
 * @author hzhlu 2023-7-13
 */
public class EMsgHelperSample {
    public static final Logger log = LoggerFactory.getLogger(EMsgHelperSample.class);

    // 原始报文对象
    static JSONObject msgObj = JSONUtil.createObj()
            .set("key1", "value1")
            .set("key2", "value2")
            .set("key3", "value3");

    public static void main(String[] args) throws NoSuchAlgorithmException, InvalidKeyException, InvalidKeySpecException {
        // 设置日志级别
//        ch.qos.logback.classic.Logger loggerImpl = (ch.qos.logback.classic.Logger) LoggerFactory.getLogger(EMsgHelper.class.getName());
//        loggerImpl.setLevel(ch.qos.logback.classic.Level.toLevel("TRACE"));

        // 显示生成的公私钥对
        // 私钥(D)：00deb3f961a5b121daaac8a2730b2830348d7f124fc211812eda27c03dc7cbc1ae
        // 公钥(Q): 04942e1fa29a58b7e10fd8eac84c1b7d1942663d7e5bc5219ffaba642119ab78e2a7df47bfeafa38a9e5ca789a773f438e79bdbe9b0d371c49caba0d463e7ba209

        // 生成密钥对并妥善保存私钥
        System.out.println("===============================================");
        System.out.println("== 生成 SM2 公私钥 ");
        KeyPair senderKeyPair = EMsgHelperV2.generateKeyPair();
        Map<String, String> senderKeyPairHex = new EMsgHelperV2().parseKeyPair(senderKeyPair);
        System.out.println("sender 私钥(D)：" + senderKeyPairHex.get("priv"));
        System.out.println("sender 公钥(Q)：" + senderKeyPairHex.get("pub"));

        KeyPair receiverKeyPair = EMsgHelperV2.generateKeyPair();
        Map<String, String> receiverKeyPairHex = new EMsgHelperV2().parseKeyPair(receiverKeyPair);
        System.out.println("receiver 私钥(D)：" + receiverKeyPairHex.get("priv"));
        System.out.println("receiver 公钥(Q)：" + receiverKeyPairHex.get("pub"));

        System.out.println("\n===============================================");
        String plainTxt = msgObj.toString();
        System.out.println("待加密的原始数据: " + plainTxt);

        // 使用样例1， 简单加解密
        System.out.println("\n===============================================");
        System.out.println("== 使用方法1： 不设置msgId（缺省值）使用HEX格式公私钥加解密报文。");
        usage1(senderKeyPairHex, receiverKeyPairHex, plainTxt);

        // 使用样例2, 加密报文中包含msgId
        System.out.println("\n===============================================");
        System.out.println("== 使用方法2： 设置msgId，使用HEX格式公私钥加解密报文。");
        usage2(senderKeyPairHex, receiverKeyPairHex, plainTxt);

        // 使用样例3
        System.out.println("\n===============================================");
        System.out.println("== 使用方法3： 设置msgId，使用HEX格式公私钥加解密报文。");
        usage3(senderKeyPair, receiverKeyPair, plainTxt);
    }

    private static void usage1(Map<String, String> senderKeyPairHex, Map<String, String> receiverKeyPairHex, String plainTxt) throws NoSuchAlgorithmException, InvalidKeySpecException, InvalidKeyException {
        // 步骤1： 发送者使用 接收者公钥 加密报文
        String receiverPrivKeyHex = receiverKeyPairHex.get("priv");
        String receiverPubKeyHex = receiverKeyPairHex.get("pub");

        String senderPrivKeyHex = senderKeyPairHex.get("priv");
        String senderPubKeyHex = senderKeyPairHex.get("pub");

        System.out.println("=== 生成密文 ===");
        System.out.println("== 1. 使用发送者的【私钥】签名报文, 私钥：" + senderPrivKeyHex);
        System.out.println("== 2. 使用接收者的【公钥】加密报文, 公钥：" + receiverPubKeyHex);
        EMsgHelperV2 encHelper = new EMsgHelperV2().setSenderPrivKey(senderPrivKeyHex).setReceiverPubKey(receiverPubKeyHex);
        String encTxt = encHelper.encrypt(plainTxt.getBytes(StandardCharsets.UTF_8));
        System.out.println("== 密文: " + encTxt);

        // 步骤3： 发送者使用 接收者公钥 加密报文
        System.out.println("=== 解密报文 ===");
        System.out.println("== 1. 使用接收者的【私钥】解密报文, 私钥：" + receiverPrivKeyHex);
        System.out.println("== 2. 使用发送者的【公钥】验签报文, 公钥：" + senderPubKeyHex);
        EMsgHelperV2 decHelper = new EMsgHelperV2().setReceiverPrivKey(receiverPrivKeyHex).setSenderPubKey(senderPubKeyHex);
        String plainTxt2 = new String(decHelper.decrypt(encTxt), StandardCharsets.UTF_8);
        System.out.println("== 解密后明文：: " + plainTxt2);
        Assert.isTrue(plainTxt2.equalsIgnoreCase(plainTxt));
    }

    private static void usage2(Map<String, String> senderKeyPairHex, Map<String, String> receiverKeyPairHex, String plainTxt) throws NoSuchAlgorithmException, InvalidKeySpecException, InvalidKeyException {
        // 消息标识:
        //  - 最大长度40个字符
        //  - 不能包含 “.”
        //  - msgId最终会出现在密文首字段，可用于标识消息使用的外界密钥，
        //    可通过此值定位合作伙伴的密钥，便于解密，如果不填此值，缺省为 JES
        String msgId = "A12-34￥51$2_er5#@%4中文56";
        System.out.println("== msgId用法说明： \n" +
                "  1. msgId参与报文加解密过程，可当做报文加密过程的盐\n" +
                "  2. msgId可用于和私钥持有者进行绑定标识，当做partnerId使用。例如，当partner发送来的加密报文，msgId设置为partnerId的值，接收者解密时通过partnerId获取对应公钥后进行解密处理。\n" +
                "  3. msgId长度小于等于40个字符，可包含字母、数字、符号、中文，但不能包含'.'");
        System.out.println("== msgId: " + msgId);

        String receiverPrivKeyHex = receiverKeyPairHex.get("priv");
        String receiverPubKeyHex = receiverKeyPairHex.get("pub");

        String senderPrivKeyHex = senderKeyPairHex.get("priv");
        String senderPubKeyHex = senderKeyPairHex.get("pub");

        // 步骤1： 发送者使用 接收者公钥 加密报文
        System.out.println("=== 生成密文 ===");
        System.out.println("== 1. 使用发送者的【私钥】签名报文, 私钥：" + senderPrivKeyHex);
        System.out.println("== 2. 使用接收者的【公钥】加密报文, 公钥：" + receiverPubKeyHex);
        EMsgHelperV2 encHelper = new EMsgHelperV2(msgId).setSenderPrivKey(senderPrivKeyHex).setReceiverPubKey(receiverPubKeyHex);
        String encTxt = encHelper.encrypt(plainTxt.getBytes(StandardCharsets.UTF_8));
        System.out.println("== 密文: " + encTxt);

        // 步骤3： 发送者使用 接收者公钥 加密报文
        System.out.println("=== 解密报文 ===");
        System.out.println("== 1. 使用接收者的【私钥】解密报文, 私钥：" + receiverPrivKeyHex);
        System.out.println("== 2. 使用发送者的【公钥】验签报文, 公钥：" + senderPubKeyHex);
        EMsgHelperV2 decHelper = new EMsgHelperV2(msgId).setReceiverPrivKey(receiverPrivKeyHex).setSenderPubKey(senderPubKeyHex);
        String plainTxt2 = new String(decHelper.decrypt(encTxt), StandardCharsets.UTF_8);
        System.out.println("== 解密后明文：" + plainTxt2);
        Assert.isTrue(plainTxt2.equalsIgnoreCase(plainTxt));
    }


    private static void usage3(KeyPair senderKeyPair, KeyPair receiverKeyPair, String plainTxt) throws NoSuchAlgorithmException, InvalidKeySpecException, InvalidKeyException {
        // 步骤1： 发送者使用 接收者公钥 加密报文
        System.out.println("=== 生成密文 ===");
        System.out.println("== 1. 使用发送者的【私钥】签名报文");
        System.out.println("== 2. 使用接收者的【公钥】加密报文");
        EMsgHelperV2 encHelper = new EMsgHelperV2().setReceiverPubKey(receiverKeyPair.getPublic()).setSenderPrivKey(senderKeyPair.getPrivate());
        String encTxt = encHelper.encrypt(plainTxt.getBytes(StandardCharsets.UTF_8));
        System.out.println("== 密文: " + encTxt);

        // 步骤2： 发送者使用 接收者公钥 加密报文
        System.out.println("=== 解密报文 ===");
        System.out.println("== 1. 使用接收者的【私钥】解密报文");
        System.out.println("== 2. 使用发送者的【公钥】验签报文");
        EMsgHelperV2 decHelper = new EMsgHelperV2().setReceiverPrivKey(receiverKeyPair.getPrivate()).setSenderPubKey(senderKeyPair.getPublic());
        String plainTxt2 = new String(decHelper.decrypt(encTxt), StandardCharsets.UTF_8);
        System.out.println("== 解密后明文：" + plainTxt2);
        Assert.isTrue(plainTxt2.equalsIgnoreCase(plainTxt));
    }


}
